import numpy as np
import yaml
from typing import List, Dict, Any, Optional

class DQTransform:
    """Base class for data quality transforms operating on a single sample (X, Y).
    X shape: [seq_len, num_nodes, input_dim]
    Y shape: [horizon, num_nodes, output_dim]
    """
    def __call__(self, x: np.ndarray, y: Optional[np.ndarray] = None):
        return x, y


class LinearInterpMissing(DQTransform):
    """Interpolate NaNs linearly over the temporal axis for each node/channel."""
    def __init__(self, max_gap: Optional[int] = None):
        self.max_gap = max_gap

    def __call__(self, x: np.ndarray, y: Optional[np.ndarray] = None):
        # operate on value channel (index 0) and optionally others; here apply to all numeric channels
        T, N, C = x.shape
        for n in range(N):
            for c in range(C):
                arr = x[:, n, c]
                if not np.any(np.isnan(arr)):
                    continue
                idx = np.arange(T)
                mask = ~np.isnan(arr)
                if mask.sum() == 0:
                    continue
                # optional: skip large gaps
                if self.max_gap is not None:
                    # identify contiguous segments of missing; if any gap longer than max_gap, leave them as NaN
                    # simple approach: proceed anyway; advanced gap control can be added later
                    pass
                arr[~mask] = np.interp(idx[~mask], idx[mask], arr[mask])
                x[:, n, c] = arr
        return x, y


class Winsorize(DQTransform):
    """Clip extremes to percentiles per sample (time-wise)."""
    def __init__(self, p_low: float = 1.0, p_high: float = 99.0):
        self.p_low = p_low
        self.p_high = p_high

    def __call__(self, x: np.ndarray, y: Optional[np.ndarray] = None):
        # clip per-channel using percentiles over time
        T, N, C = x.shape
        for n in range(N):
            for c in range(C):
                arr = x[:, n, c]
                try:
                    lo = np.nanpercentile(arr, self.p_low)
                    hi = np.nanpercentile(arr, self.p_high)
                    x[:, n, c] = np.clip(arr, lo, hi)
                except Exception:
                    pass
        return x, y


class ExtremeMask(DQTransform):
    """Optionally mask or zero-out extreme values beyond k*std (z-score based)."""
    def __init__(self, std_k: float = 5.0, mask_to_zero: bool = True):
        self.std_k = std_k
        self.mask_to_zero = mask_to_zero

    def __call__(self, x: np.ndarray, y: Optional[np.ndarray] = None):
        T, N, C = x.shape
        for n in range(N):
            for c in range(C):
                arr = x[:, n, c]
                mu = np.nanmean(arr)
                sd = np.nanstd(arr)
                if sd <= 1e-12:
                    continue
                z = (arr - mu) / sd
                mask = np.abs(z) > self.std_k
                if self.mask_to_zero:
                    arr[mask] = 0.0
                else:
                    arr[mask] = mu
                x[:, n, c] = arr
        return x, y


def build_dq_transforms(config_path: Optional[str]) -> List[DQTransform]:
    """Load DQ pipeline from YAML config or return sensible defaults."""
    if not config_path:
        # defaults: interpolate + winsorize
        return [LinearInterpMissing(), Winsorize(p_low=1.0, p_high=99.0)]
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            cfg: Dict[str, Any] = yaml.safe_load(f) or {}
    except Exception:
        return [LinearInterpMissing(), Winsorize(p_low=1.0, p_high=99.0)]
    out: List[DQTransform] = []
    dq_items = cfg.get('transforms', [])
    for item in dq_items:
        name = str(item.get('name', '')).lower()
        if name == 'interpolate_missing':
            out.append(LinearInterpMissing(max_gap=item.get('max_gap', None)))
        elif name == 'winsorize':
            out.append(Winsorize(p_low=float(item.get('p_low', 1.0)), p_high=float(item.get('p_high', 99.0))))
        elif name == 'extreme_mask':
            out.append(ExtremeMask(std_k=float(item.get('std_k', 5.0)), mask_to_zero=bool(item.get('mask_to_zero', True))))
    if not out:
        out = [LinearInterpMissing(), Winsorize(p_low=1.0, p_high=99.0)]
    return out